

package com.hazelcast.jet.pipeline.test;

import com.hazelcast.jet.Traverser;
import com.hazelcast.jet.Traversers;
import com.hazelcast.jet.core.AbstractProcessor;

import javax.annotation.Nonnull;
import java.util.List;
import java.util.stream.IntStream;

import static java.util.stream.Collectors.toList;


/**
 * Implements a mock batch source with parallelism. It emits the items from
 * the lists you supply, preserving the order within each list.
 *
 * @param <T> type of emitted items
 */
public class ParallelBatchP<T> extends AbstractProcessor {

    private final List<? extends Iterable<T>> iterables;
    private List<Traverser<T>> traversers;

    /**
     * Creates a processor that generates items using its assigned
     * iterables. This processor picks its assigned iterables from
     * the list of iterables according to its global processor index.
     *
     * @param iterables list of iterables used to emit items
     */
    public ParallelBatchP(List<? extends Iterable<T>> iterables) {
        this.iterables = iterables;
    }

    @Override
    protected void init(@Nonnull Context context) {
        int globalProcessorIndex = context.globalProcessorIndex();
        int totalParallelism = context.totalParallelism();
        traversers = IntStream.range(0, iterables.size()).filter(i -> i % totalParallelism == globalProcessorIndex).mapToObj(iterables::get).map(Traversers::traverseIterable).collect(toList());
    }

    @Override
    public boolean complete() {
        boolean isCompleted = true;
        for (Traverser<T> traverser : traversers) {
            if (!emitFromTraverser(traverser)) {
                isCompleted = false;
            }
        }
        return isCompleted;
    }
}
